# ==============================================================================
# analysis-main.R
# author: Anselm Hager / Bernhard Clemm
# ==============================================================================

dir <- dirname(dirname(rstudioapi::getSourceEditorContext()$path))
source(paste0(dir, "/code/setup-packages.R"))

# DATA =========================================================================

load(paste0(dir, "/data/data_processed.RData"))

# TWFE ESTIMATOR ===============================================================

run_twfe <- function(dat, controls = NULL) {
  twfe_results <- vector(mode = "list", length = 3)

  if (is.null(controls)) {
    form <- as.formula(
      "percent ~ treat_cont | ags + election_date_int"
    )
  } else {
    form <- as.formula(
      paste("percent ~ treat_cont + ", paste(controls, collapse = " + "), sep = "") %>%
        paste(., "| ags + election_date_int", sep = " ")
    )
  }

  for (i in 1:length(twfe_results)) {
    party_results <- vector(mode = "list", length = 2)
    parties <- c("afd", "green")

    for (j in 1:length(party_results)) {
      twfe <- fixest::feols(form, data = dplyr::filter(dat[[i]], party == parties[j]),
                            collin.tol = 1e-30) %>%
        broom::tidy(conf.int = TRUE) %>%
        dplyr::filter(grepl("treat_cont", term)) %>%
        dplyr::mutate(
          method = "TWFE",
          party = parties[j]
        )

      party_results[[j]] <- twfe
    }

    twfe_results[[i]] <- dplyr::bind_rows(party_results)
  }

  rm(twfe, i, j, parties, party_results)
  names(twfe_results) <- names(dat)
  return(twfe_results)
}

twfe <- run_twfe(data_processed)
twfe_controls <- run_twfe(
  data_processed,
  controls = c("prop_agri", "prop_heathen", "prop_unempl", "refugees")
)

twfe_lags <- run_twfe(data_processed, controls = c("treat_cont_lag_1"))
twfe_lags_controls <- run_twfe(
  data_processed,
  controls = c("prop_agri", "prop_heathen", "prop_unempl", "refugees", "treat_cont_lag_1")
)

# PANELMATCH ESTIMATOR (IMAI ET AL) ============================================

run_panel_match <- function(dat, controls = NULL) {
  panel_match_results <- vector(mode = "list", length = 3)

  for (i in 1:length(panel_match_results)) {
    panel_match_party <- vector(mode = "list", length = 2)
    parties <- c("afd", "green")

    for (j in 1:length(parties)) {
      dat_run <- as.data.frame(dplyr::filter(dat[[i]], party == parties[j])) %>%
        dplyr::arrange(ags) %>%
        dplyr::mutate(
          election_date_int = as.integer(election_date_int),
          ags_int = as.numeric(as.factor(ags)) * 100
        ) %>%
        dplyr::ungroup()


      if (is.null(controls)) {
        message(paste("constructing matched panel", i, "for", parties[j], sep = " "))

        pm_est <- try(
          PanelMatch::PanelMatch(
            lag = 5,
            time.id = "election_date_int",
            unit.id = "ags_int",
            treatment = "attacks_binary",
            refinement.method = "none",
            data = dat_run,
            match.missing = TRUE,
            size.match = 5,
            qoi = "att",
            outcome.var = "percent",
            forbid.treatment.reversal = FALSE,
            use.diagonal.variance.matrix = TRUE
          )
        )
      } else {
        form <- paste(controls, collapse = " + ") %>%
          paste0("~ ", .) %>%
          as.formula()

        message(paste("constructing matched panel", i, "for", parties[j], sep = " "))

        pm_est <- try(
          PanelMatch::PanelMatch(
            lag = 5,
            time.id = "election_date_int",
            unit.id = "ags_int",
            treatment = "attacks_binary",
            refinement.method = "mahalanobis",
            data = dat_run,
            covs.formula = form,
            match.missing = FALSE,
            size.match = 5,
            qoi = "att",
            outcome.var = "percent",
            forbid.treatment.reversal = FALSE,
            use.diagonal.variance.matrix = TRUE
          )
        )
      }

      message(paste("computing panel match estimator", i, "for", parties[j], sep = " "))

      if (!is(pm_est, "try-error")) {
        pm_est <- PanelMatch::PanelEstimate(sets = pm_est, data = dat_run)

        panel_match_party[[j]] <- tibble(
          t = c(0),
          estimate = summary(pm_est)$summary[, 1],
          conf.low = summary(pm_est)$summary[, 3],
          conf.high = summary(pm_est)$summary[, 4]
        ) %>%
          dplyr::select(t, estimate, conf.low, conf.high) %>%
          dplyr::mutate(
            method = "Panel Match",
            party = parties[j]
          ) %>%
          as.data.frame()
      } else {
        panel_match_party[[j]] <- data.frame(
          t = NA,
          estimate = NA,
          conf.low = NA,
          conf.high = NA,
          method = "Panel Match",
          party = parties[j]
        )
      }
    }

    panel_match_results[[i]] <- dplyr::bind_rows(panel_match_party)
  }

  return(panel_match_results)
}

panel_match <- run_panel_match(data_processed)
names(panel_match) <- names(data_processed)

panel_match_controls <- run_panel_match(
  data_processed,
  controls = c("prop_agri", "prop_heathen", "prop_unempl", "refugees")
)
names(panel_match_controls) <- names(data_processed)

# EFFECT PLOT (Figure 2) =======================================================

plot_res <- list(
  twfe = twfe,
  twfe_controls = twfe_controls,
  twfe_lags = twfe_lags,
  twfe_lags_controls = twfe_lags_controls,
  panel_match = panel_match,
  panel_match_controls = panel_match_controls
)

## Remove lagged treatment coefficient (only used as control) ####
plot_res$twfe_lags <- lapply(plot_res$twfe_lags, function(x) {
  x %>% filter(term != "treat_cont_lag_1")
})
plot_res$twfe_lags_controls <- lapply(plot_res$twfe_lags_controls, function(x) {
  x %>% filter(term != "treat_cont_lag_1")
})

## Select plotting variables ####
for (i in 1:length(plot_res)) {
  plot_res[[i]] <- lapply(plot_res[[i]], function(x) {
    x %>% dplyr::select(c(estimate, conf.low, conf.high, method, party))
  })
}

## Add election variable ####
for (i in 1:length(plot_res)) {
  election <- names(plot_res[[i]])

  for (j in 1:3) {
    plot_res[[i]][[j]] <- plot_res[[i]][[j]] %>%
      dplyr::mutate(election = election[j])
  }
}
rm(election)

## Combine election estimates for each mode ####
for (i in 1:length(plot_res)) {
  plot_res[[i]] <- dplyr::bind_rows(plot_res[[i]])
}

## Give lagged treatment TWFE model appropriate name ####
plot_res$twfe_lags <- plot_res$twfe_lags %>%
  dplyr::mutate(method = "TWFE lagged treatment")
plot_res$twfe_lags_controls <- plot_res$twfe_lags_controls %>%
  dplyr::mutate(method = "TWFE lagged treatment")

## Add controls variable indicator ####
cont <- which(grepl("_controls\\b", names(plot_res)))

for (i in cont) {
  plot_res[[i]] <- plot_res[[i]] %>%
    dplyr::mutate(controls = 1)
}

for (i in c(1:length(plot_res))[-cont]) {
  plot_res[[i]] <- plot_res[[i]] %>%
    dplyr::mutate(controls = 0)
}

plot_res <- dplyr::bind_rows(plot_res)

plot_res <- plot_res %>%
  dplyr::mutate(election_pos = dplyr::case_when(
    election == "btw" ~ 36,
    election == "ltw" ~ 24,
    election == "kw" ~ 12
  )) %>%
  dplyr::mutate(method_pos = dplyr::case_when(
    method == "TWFE" ~ 3,
    method == "TWFE lagged treatment" ~ 0,
    method == "Panel Match" ~ -3
  )) %>%
  dplyr::mutate(party_pos = dplyr::case_when(
    party == "green" ~ 0.55,
    party == "afd" ~ -0.55
  )) %>%
  dplyr::mutate(control_pos = dplyr::case_when(
    controls == 1 ~ 0.25,
    controls == 0 ~ -0.25
  )) %>%
  dplyr::mutate(pos = election_pos + method_pos + party_pos + control_pos) %>%
  dplyr::mutate(
    controls = ifelse(controls == 1, "treatment + controls", "treatment only"),
    party = ifelse(party == "afd", "AfD", "Greens"),
    method = ifelse(method == "Panel Match", "PanelMatch", method)
  ) %>%
  dplyr::mutate(
    controls = factor(controls,
      ordered = TRUE,
      levels = c("treatment + controls", "treatment only")
    ),
    method = factor(method,
      ordered = TRUE,
      levels = c("TWFE", "TWFE lagged treatment", "PanelMatch")
    )
  ) %>%
  dplyr::mutate(
    estimate = estimate * 100,
    conf.low = conf.low * 100,
    conf.high = conf.high * 100
  ) %>%
  dplyr::arrange(plot_res, pos)

plot_res %>%
  ggplot(aes(
    x = estimate, y = pos, color = party,
    shape = method, linetype = controls
  )) +
  geom_point(size = 2.2) +
  geom_errorbarh(aes(xmax = conf.high, xmin = conf.low),
    height = 0.5, size = 0.5
  ) +
  scale_y_continuous(
    breaks = c(12, 24, 36),
    labels = c(
      "Wolf attack \n (local \n elections)",
      "Wolf attack \n (state \n elections)",
      "Wolf attack \n (federal \n elections)"
    )
  ) +
  scale_x_continuous(name = "Effect (percentage points)") +
  scale_color_manual(values = c("blue2", "green3"), name = "Party") +
  scale_shape(name = "Model") +
  scale_linetype(name = "") +
  geom_vline(xintercept = 0, linetype = "dashed", color = "black") +
  theme_bw() +
  theme(
    legend.position = "bottom",
    legend.title = element_text(size = 14),
    legend.text = element_text(size = 10),
    axis.title.y = element_blank(),
    axis.text.y = element_text(size = 14),
    axis.text.x = element_text(size = 14),
    axis.ticks.y = element_blank(),
    panel.grid.major.y = element_blank(),
    legend.background = element_blank(),
    legend.box.background = element_rect(colour = "black"),
    text = element_text(size = 14)
  ) +
  guides(
    color = guide_legend(nrow = 2, order = 1),
    shape = guide_legend(nrow = 3, order = 2),
    linetype = guide_legend(nrow = 2, order = 3)
  )

# ggsave(paste0(dir, "/pretrends_btw.pdf"))
